import cv2
import numpy as np
import tensorflow as tf
import json
import os
from logger import Logger

class ImagePredictor:
    def __init__(self, model_path, class_names_path, img_size=224):
        self.logger = Logger(name='predictor')
        self.img_size = img_size
        self.model = tf.keras.models.load_model(model_path)
        self.class_names = self._load_class_names(class_names_path)
        
    def _load_class_names(self, class_names_path):
        """加载类别名称"""
        with open(class_names_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def preprocess_image(self, image_path, return_original=False):
        """预处理图片，可选择返回原图"""
        # 读取原始图片
        original_img = cv2.imread(image_path)
        if original_img is None:
            raise ValueError(f"无法读取图片: {image_path}")
            
        # 保存原始尺寸
        self.original_height, self.original_width = original_img.shape[:2]
        
        # 预处理图片
        img = cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB)
        processed_img = cv2.resize(img, (self.img_size, self.img_size))
        processed_img = processed_img.astype(np.float32) / 255.0
        processed_img = np.expand_dims(processed_img, axis=0)
        
        if return_original:
            return processed_img, original_img
        return processed_img
    
    def predict_and_draw(self, image_path, confidence_threshold=0.5):
        """预测并在图片上绘制标注"""
        try:
            # 预处理图片，同时获取原图
            img_processed, original_img = self.preprocess_image(image_path, return_original=True)
            
            # 获取预测结果
            predictions = self.model.predict(img_processed)
            class_idx = np.argmax(predictions[0])
            confidence = float(predictions[0][class_idx])
            
            # 如果置信度超过阈值，绘制标注
            if confidence >= confidence_threshold:
                # 获取预测的类别名称
                class_name = self.class_names[class_idx]
                
                # 在整个图片范围内绘制矩形
                # 为了留出文字空间，顶部留出一些边距
                text_margin = 40
                cv2.rectangle(original_img, 
                            (5, text_margin),  # 左上角坐标
                            (self.original_width-5, self.original_height-5),  # 右下角坐标
                            (0, 255, 0),  # 绿色边框
                            2)  # 线条粗细
                
                # 添加文字标签
                label = f"{class_name}: {confidence:.2%}"
                cv2.putText(original_img,
                           label,
                           (10, 30),  # 文字位置
                           cv2.FONT_HERSHEY_SIMPLEX,  # 字体
                           0.8,  # 字体大小
                           (0, 255, 0),  # 绿色文字
                           2)  # 线条粗细
                
                # 保存标注后的图片
                output_path = os.path.join(
                    'predictions',
                    f"pred_{os.path.basename(image_path)}"
                )
                os.makedirs('predictions', exist_ok=True)
                cv2.imwrite(output_path, original_img)
                
                self.logger.info(f"预测结果已保存至: {output_path}")
                
                return {
                    'class_name': class_name,
                    'confidence': confidence,
                    'output_path': output_path
                }
            else:
                self.logger.info("未检测到置信度足够高的物体")
                return None
                
        except Exception as e:
            self.logger.error(f"预测过程出错: {str(e)}")
            raise
    
    def predict_batch(self, image_dir, confidence_threshold=0.5):
        """批量预测文件夹中的图片"""
        results = []
        for image_name in os.listdir(image_dir):
            if image_name.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_path = os.path.join(image_dir, image_name)
                result = self.predict_and_draw(image_path, confidence_threshold)
                if result:
                    results.append(result)
        return results